##########################################################################################

library(data.table)
library(optparse)
library(pheatmap)
library(GSVA)

##########################################################################################

option_list <- list(
    make_option(c("--data_type"), type = "character"),
    make_option(c("--rna_file"), type = "character"),
    make_option(c("--cor_file"), type = "character"),
    make_option(c("--clust"), type = "character"),
    make_option(c("--mfuzz_file"), type = "character"),
    make_option(c("--cor_t"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    rna_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/Motif.MeanByCellType.tsv"
    cor_file <- "~/20231121_singleMuti/results/celltype_plot/mfuzz/cor.motif_atac-rna.tsv"
    data_type <- "motif"
    clust <- 5
    mfuzz_file <- "~/20231121_singleMuti/results/celltype_plot/mfuzz/mfuzz_plot.motif.5.tsv"
    out_path <- "~/20231121_singleMuti/results/celltype_plot/mfuzz"
    cor_t <- 0.5
}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

rna_file <- opt$rna_file
cor_file <- opt$cor_file
data_type <- opt$data_type
mfuzz_file <- opt$mfuzz_file
cor_t <- as.numeric(opt$cor_t)
out_path <- opt$out_path
clust <- as.numeric(opt$clust)

dir.create(out_path , recursive = T)

##########################################################################################

dat_mfuzz <- data.frame(fread(mfuzz_file))
dat_cor <- data.frame(fread(cor_file))
sco_exp <- data.frame(fread(rna_file))

set.seed(1234)

##########################################################################################

cell_order <- c("SSC" , "Differenting.Differented.SPG" , "Leptotene" ,
    "Zygotene" , "Patchytene" , "Diplotene" , 
    "Early.stage.of.spermatids" , "Round.ElongateS.tids" , "Sperm"
    )

##########################################################################################
## 基因表达矩阵
rownames(sco_exp) <- sco_exp$gene
sco_exp <- sco_exp[,-1]

if(data_type == "exp"){
    sco_exp <- sco_exp[which(rowSums(sco_exp)>0),]
}
sco_exp <- sco_exp[,cell_order]
sco_exp <- t(sco_exp)

##########################################################################################

dat_mfuzz <- merge( dat_mfuzz , dat_cor , by.x = "NAME" , by.y = "GeneExpressionMatrix_name" , all.x = T )

## 阳性转录因子
postive_tf <- subset(dat_cor , cor > cor_t)$GeneExpressionMatrix_name
other_tf <- subset(dat_cor , cor < cor_t)$GeneExpressionMatrix_name

## 计算富集
dat_mfuzz$TF <- ifelse( dat_mfuzz$NAME %in% dat_cor$GeneExpressionMatrix_name , "TRUE" , "FALSE" )
dat_mfuzz$postive_tf <- ""
dat_mfuzz$postive_tf <- ifelse( dat_mfuzz$NAME %in% postive_tf , "positive" , dat_mfuzz$postive_tf )
dat_mfuzz$postive_tf <- ifelse( dat_mfuzz$NAME %in% other_tf , "other" , dat_mfuzz$postive_tf )

out_name <- paste0(out_path , "/mfuzz_plot." , data_type , "." , clust , ".motif_clust.motifAnnotation.tsv")
write.table( dat_mfuzz , out_name , row.names = F , quote = F , sep = "\t" )


tmp_dat <- table(dat_mfuzz$postive_tf , dat_mfuzz$CLUSTER)[c("positive" , "other"),]
## 分每个cluster计算富集
result <- c()
for(clusN in colnames(tmp_dat)){
    a <- tmp_dat[,clusN]
    b <- apply(tmp_dat[,colnames(tmp_dat)!=clusN] , 1 , sum)
    tmp_fisher <- fisher.test(data.frame(a,b))

    p <- tmp_fisher$p.value
    er <- tmp_fisher$estimate
    L95 <- tmp_fisher$conf.int[1] 
    U95 <- tmp_fisher$conf.int[2]

    tmp_res <- data.frame( clusN = clusN , P = p , ER = er , L95 = L95 , U95 = U95  )
    result <- rbind(result , tmp_res)
}

out_name <- paste0(out_path , "/mfuzz_plot." , data_type , "." , clust , ".motif_enrich." , cor_t , ".tsv")
write.table( result , out_name , row.names = F , quote = F , sep = "\t" )

out_name <- paste0(out_path , "/mfuzz_plot." , data_type , "." , clust , ".motif_clust." , cor_t , ".tsv")
write.table( tmp_dat , out_name , row.names = T , quote = F , sep = "\t" )

##########################################################################################
## 按照基因的cluster进行gsva
gsva_path <- data.frame( gene = dat_mfuzz$NAME , group_name = dat_mfuzz$CLUSTER )

tmp <- list()
for (i in unique(gsva_path$group_name)){
  i <- list(subset(gsva_path,group_name==i)$gene)
  tmp <- append(tmp,i)
}

names(tmp) <- unique(gsva_path$group_name)
gsva_score <- gsva(as.matrix(t(sco_exp)),tmp,method="ssgsea",verbose=T,parallel.sz=40)

out_file <- paste0( out_path , "/mfuzz_plot." , data_type , "." , clust , ".gsva.tsv" )
write.table(gsva_score, out_file , row.names = T , quote = F , sep = "\t")

p <- pheatmap(t(gsva_score),
scale = "column",
show_rownames = T ,show_colnames = T, 
#cutree_cols=7,cutree_rows = 7,
cluster_rows = F , cluster_cols = T , clustering_method = "ward.D2",
color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
cellheight = 10)
out_file <- paste0( out_path , "/mfuzz_plot." , data_type , "." , clust , ".gsva.pdf" )
pdf(out_file)
print(p)
dev.off()